
from exp import data, run
from util import color
import random, os
from datetime import datetime
from . import k_shot

class Kwait_Reasoner:
	def __init__(self, examples, k_shot=8, k=0):
		self.examples = examples
		self.k_shot = k_shot
		self.k = k

	def reason(self, model_id, question, verbal=True):
		random.shuffle(self.examples)
		examples = [e for e in self.examples[:self.k_shot+1] if e['problem'] != question]
		examples = k_shot._examples_repr(examples[:self.k_shot])
		partial_solution = ''
		for i in range(self.k + 1):
			if i > 0:
				partial_solution = partial_solution + '\nWait,'
			resp = run.generate(model_id, question, 'com/reasoner_prompts', 'k_wait',
								question=question.question, 
								examples=examples, 
								verbal=verbal,
								partial_solution=partial_solution,
								)
			partial_solution = partial_solution + resp
		return question.extra_asnwer(resp), 'The answer is', []

def test(test_set):
	math500 = data.get_data('MATH500')
	test_examples = data.get_data(test_set)
	reasoner = Kwait_Reasoner(math500, k=2)
	for e in test_examples[:1]:
		run.test(e, reasoner.reason)

out_path = '../OUT/k_shot'

def exp(test_set):
	math500 = data.get_data('MATH500')
	test_examples = data.get_data(test_set)
	for k in [1, 2, 3]:
		reasoner = Kwait_Reasoner(math500, k=k)
		outfile = f'{out_path}/{test_set}_{k}-wait.jsonl'
		print(f'# {k} {outfile} {datetime.now()}')
		if os.path.isfile(outfile): continue
		run.run_exp(test_examples, reasoner.reason, num_samples=1, max_retries=8, outfile=outfile)

def main():
	if not os.path.isdir(out_path):
		os.mkdir(out_path)
	exp('MATH500')
	exp('AIME')

if __name__ == '__main__':
	main()
